// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/aggregates/average_aggregate.h>
#include <pollux/functions/lib/aggregates/average_aggregate_base.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>

using namespace kumo::pollux::functions::aggregate;

namespace kumo::pollux::aggregate::prestosql {
    /// Count is BIGINT() while sum and the final aggregates type depends on
    /// the input types:
    ///       INPUT TYPE    |     SUM             |     AVG
    ///   ------------------|---------------------|------------------
    ///     DOUBLE          |     DOUBLE          |    DOUBLE
    ///     REAL            |     DOUBLE          |    REAL
    ///     ALL INTs        |     DOUBLE          |    DOUBLE
    ///     DECIMAL         |     DECIMAL         |    DECIMAL
    void registerAverageAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;

        for (const auto &inputType: {"smallint", "integer", "bigint", "double"}) {
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .returnType("double")
                .intermediateType("row(double,bigint)")
                .argumentType(inputType)
                .build());
        }

        // Interval input type in Presto has special case and returns INTERVAL, not
        // DOUBLE.
        signatures.push_back(exec::AggregateFunctionSignatureBuilder()
            .returnType("interval day to second")
            .intermediateType("row(double,bigint)")
            .argumentType("interval day to second")
            .build());

        // Real input type in Presto has special case and returns REAL, not DOUBLE.
        signatures.push_back(exec::AggregateFunctionSignatureBuilder()
            .returnType("real")
            .intermediateType("row(double,bigint)")
            .argumentType("real")
            .build());

        signatures.push_back(exec::AggregateFunctionSignatureBuilder()
            .integerVariable("a_precision")
            .integerVariable("a_scale")
            .argumentType("DECIMAL(a_precision, a_scale)")
            .intermediateType("varbinary")
            .returnType("DECIMAL(a_precision, a_scale)")
            .build());

        auto name = prefix + kAvg;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_LE(
                    argTypes.size(), 1, "{} takes at most one argument", name);
                auto inputType = argTypes[0];
                if (exec::isRawInput(step)) {
                    switch (inputType->kind()) {
                        case TypeKind::SMALLINT:
                            return std::make_unique<
                                AverageAggregateBase<int16_t, double, double> >(resultType);
                        case TypeKind::INTEGER:
                            return std::make_unique<
                                AverageAggregateBase<int32_t, double, double> >(resultType);
                        case TypeKind::BIGINT: {
                            if (inputType->isShortDecimal()) {
                                return std::make_unique<DecimalAverageAggregateBase<int64_t> >(
                                    resultType);
                            }
                            if (inputType->isIntervalDayTime()) {
                                return std::make_unique<
                                    AverageAggregateBase<int64_t, double, int64_t> >(resultType);
                            }
                            return std::make_unique<
                                AverageAggregateBase<int64_t, double, double> >(resultType);
                        }
                        case TypeKind::HUGEINT: {
                            if (inputType->isLongDecimal()) {
                                return std::make_unique<DecimalAverageAggregateBase<int128_t> >(
                                    resultType);
                            }
                            POLLUX_NYI();
                        }
                        case TypeKind::REAL:
                            return std::make_unique<
                                AverageAggregateBase<float, double, float> >(resultType);
                        case TypeKind::DOUBLE:
                            return std::make_unique<
                                AverageAggregateBase<double, double, double> >(resultType);
                        default:
                            POLLUX_FAIL(
                                "Unknown input type for {} aggregation {}",
                                name,
                                inputType->kindName());
                    }
                } else {
                    checkAvgIntermediateType(inputType);
                    switch (resultType->kind()) {
                        case TypeKind::REAL:
                            return std::make_unique<
                                AverageAggregateBase<int64_t, double, float> >(resultType);
                        case TypeKind::DOUBLE:
                        case TypeKind::ROW:
                            return std::make_unique<
                                AverageAggregateBase<int64_t, double, double> >(resultType);
                        case TypeKind::BIGINT:
                            return std::make_unique<DecimalAverageAggregateBase<int64_t> >(
                                resultType);
                        case TypeKind::HUGEINT:
                            return std::make_unique<DecimalAverageAggregateBase<int128_t> >(
                                resultType);
                        case TypeKind::VARBINARY:
                            if (inputType->isLongDecimal()) {
                                return std::make_unique<DecimalAverageAggregateBase<int128_t> >(
                                    resultType);
                            } else if (
                                inputType->isShortDecimal() ||
                                inputType->kind() == TypeKind::VARBINARY) {
                                // If the input and out type are VARBINARY, then the
                                // LongDecimalWithOverflowState is used and the template type
                                // does not matter.
                                return std::make_unique<DecimalAverageAggregateBase<int64_t> >(
                                    resultType);
                            }
                            [[fallthrough]];
                        default:
                            POLLUX_FAIL(
                                "Unsupported result type for final aggregation: {}",
                                resultType->kindName());
                    }
                }
            },
            {false /*orderSensitive*/, false /*companionFunction*/},
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
